import pandas as pd
import matplotlib.pyplot as plt
from sklearn.ensemble import RandomForestClassifier
import warnings

warnings.filterwarnings("ignore")
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

df = pd.read_csv('../data/train.csv')
# 取出无关列
df.drop(['EmployeeNumber', 'StandardHours', 'Over18'], axis=1, inplace=True)

# 类别列one-hot编码
categorical_cols = df.select_dtypes(include=['object']).columns.tolist()
df_encoded = pd.get_dummies(df, columns=categorical_cols, drop_first=True)

x = df_encoded.drop('Attrition', axis=1)
y = df_encoded['Attrition']

rf = RandomForestClassifier(n_estimators=100, random_state=42, n_jobs=-1)
rf.fit(x, y)

importance = pd.Series(rf.feature_importances_, index=x.columns).sort_values(ascending=False)
print(importance)

importance_df = importance.reset_index()
importance_df.columns = ['Feature', 'Importance']
importance_df.to_csv('../data/rf_FeaturesImportance.csv', index=False)

# 可视化
plt.figure(figsize=(20, 10))
# 根据重要性值创建颜色渐变
colors = plt.cm.Reds(importance.values / importance.values.max())
plt.barh(importance.index, importance.values, color=colors)
plt.title("随机森林特征重要性")
plt.xlabel("重要性")
plt.gca().invert_yaxis()  # 保持重要性从高到低的顺序
plt.tight_layout()
plt.grid(axis='x', alpha=0.7)
plt.savefig("../data/rf_FeaturesImportance.png", dpi=300)
plt.show()
